
from langchain_community.embeddings import DashScopeEmbeddings

from langchain_redis import RedisConfig, RedisVectorStore

from app.core.config import settings


class KnowledgeBaseManager:
    """管理知识库retriever"""

    def __init__(self,index_name: str,k: int = 4):
        self.redis_url = settings.REDIS_URL
        self.k = k
        self.index_name = index_name
        self.embedding_model = None
        self.vector_store = None
        self.retriever = None

        # 初始化向量存储和检索器
        self.initialize_vector_store()




    def initialize_vector_store(self):
        """初始化向量存储和检索器"""
        # 创建Embedding模型
        self.embedding_model = DashScopeEmbeddings(
            model=settings.EMBEDDING_MODEL,
            dashscope_api_key=settings.DASHSCOPE_API_KEY,
            )
        # 配置Redis
        config = RedisConfig(
            redis_url=self.redis_url,
            index_name=self.index_name,
            from_existing=True
        )

        # 初始化向量存储
        self.vector_store = RedisVectorStore(self.embedding_model, config=config)
        # 获取检索器
        self.retriever = self.vector_store.as_retriever(search_kwargs={"k": self.k})



